"""
Model
"""
import mindspore
import mindspore.nn as nn


class MLP(nn.Cell):
    """
    多层感知机(MLP)模型
    适用于2D湍流特征的学习与预测
    """
    def __init__(self, config):
        """
        初始化MLP网络
        
        Args:
            config (dict): 网络配置字典，需包含：
                - input_num: 输入特征维度
                - layer1_output-layer5_output: 各隐藏层输出维度
                - activation_function: 激活函数类型
        """
        super(MLP, self).__init__()
        # 定义5层全连接层
        self.fc1 = nn.Dense(config["input_num"], config["layer1_output"])
        self.fc2 = nn.Dense(config["layer1_output"], config["layer2_output"])
        self.fc3 = nn.Dense(config["layer2_output"], config["layer3_output"])
        self.fc4 = nn.Dense(config["layer3_output"], config["layer4_output"])
        self.fc5 = nn.Dense(config["layer4_output"], config["layer5_output"])
        
        # 根据配置选择激活函数
        if config["activation_function"] == "ReLU":
            self.actfunc = nn.ReLU()
        elif config["activation_function"] == "LeakyReLU":
            self.actfunc = nn.LeakyReLU()
        elif config["activation_function"] == "Tanh":
            self.actfunc = nn.Tanh()
        elif config["activation_function"] == "Sigmoid":
            self.actfunc = nn.Sigmoid()
        else:
            self.actfunc = nn.ReLU()  # 默认使用ReLU

    def construct(self, input_tensor):
        """
        前向传播过程
        
        Args:
            input_tensor (Tensor): 输入特征张量，形状为(batch_size, input_num)
            
        Returns:
            Tensor: 输出张量，形状为(batch_size, layer5_output)
        """
        # 各层前向传播 + 激活函数
        input_tensor = self.actfunc(self.fc1(input_tensor))
        input_tensor = self.actfunc(self.fc2(input_tensor))
        input_tensor = self.actfunc(self.fc3(input_tensor))
        input_tensor = self.actfunc(self.fc4(input_tensor))
        output = self.fc5(input_tensor)
        return output


class ResMLP(nn.Cell):
    """
    残差多层感知机(ResMLP)模型
    通过残差连接解决深层网络训练困难问题，适用于3D湍流复杂特征提取
    """
    def __init__(self, input_num, width, depth, output_num):
        """
        初始化ResMLP网络
        
        Args:
            input_num (int): 输入特征维度
            width (int): 隐藏层宽度（各残差层维度）
            depth (int): 残差层数量
            output_num (int): 输出维度（如湍流粘度预测）
        """
        super(ResMLP, self).__init__()
        # 输入映射层
        self.linear_input = nn.Dense(input_num, width)
        # 残差层列表（全连接层）
        self.layers = nn.CellList([nn.Dense(width, width) for _ in range(depth)])
        # 输出映射层
        self.linear_out = nn.Dense(width, output_num)
        # 激活函数
        self.act_func = nn.ReLU()
        # 残差连接权重参数（可训练）
        self.resw = mindspore.Parameter(mindspore.numpy.zeros(depth), requires_grad=True)

    def construct(self, input_tensor):
        """
        带残差连接的前向传播过程
        
        Args:
            input_tensor (Tensor): 输入特征张量，形状为(batch_size, input_num)
            
        Returns:
            Tensor: 输出张量，形状为(batch_size, output_num)
        """
        # 输入映射与激活
        input_tensor = self.act_func(self.linear_input(input_tensor))
        # 残差层前向传播
        for i, layer in enumerate(self.layers):
            # 残差连接公式: x = x + w * act(layer(x))
            input_tensor = input_tensor + self.resw[i] * self.act_func(layer(input_tensor))
        # 输出映射
        input_tensor = self.linear_out(input_tensor)
        return input_tensor